#!/usr/bin/env python3
"""
训练进度实时监控脚本
"""

import os
import time
import subprocess
import signal
import sys
from pathlib import Path
import logging

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class TrainingMonitor:
    """训练监控器"""
    
    def __init__(self, project_root: str = None):
        self.project_root = Path(project_root) if project_root else Path(__file__).parent.parent
        self.log_dir = self.project_root / "logs" / "universal_document"
        self.running = True
        
        # 注册信号处理
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
    
    def _signal_handler(self, signum, frame):
        """信号处理器"""
        logger.info("🛑 收到中断信号，正在停止监控...")
        self.running = False
    
    def monitor_gpu_usage(self):
        """监控GPU使用情况"""
        try:
            result = subprocess.run([
                'nvidia-smi', 
                '--query-gpu=index,name,memory.used,memory.total,utilization.gpu',
                '--format=csv,noheader,nounits'
            ], capture_output=True, text=True, timeout=5)
            
            if result.returncode == 0:
                print(f"\n{'='*60}")
                print("🖥️  GPU 状态:")
                for line in result.stdout.strip().split('\n'):
                    if line.strip():
                        parts = [p.strip() for p in line.split(',')]
                        if len(parts) >= 5:
                            gpu_id, name, used, total, util = parts
                            memory_percent = (int(used) / int(total)) * 100
                            print(f"  GPU {gpu_id}: {name}")
                            print(f"    内存: {used}MB / {total}MB ({memory_percent:.1f}%)")
                            print(f"    利用率: {util}%")
                print(f"{'='*60}\n")
        except Exception as e:
            logger.debug(f"获取GPU信息失败: {e}")
    
    def monitor_training_logs(self, stage: int = None):
        """监控训练日志"""
        if stage:
            log_files = list(self.log_dir.glob(f"stage{stage}/*_training.log"))
        else:
            log_files = list(self.log_dir.glob("stage*/*_training.log"))
        
        if not log_files:
            logger.warning("🔍 未找到训练日志文件")
            return
        
        # 使用最新的日志文件
        latest_log = max(log_files, key=lambda x: x.stat().st_mtime)
        
        logger.info(f"📝 监控日志文件: {latest_log}")
        
        # 实时监控日志
        try:
            with open(latest_log, 'r', encoding='utf-8') as f:
                # 先读取已有内容的最后几行
                f.seek(0, 2)  # 移动到文件末尾
                file_size = f.tell()
                
                # 如果文件不为空，显示最后几行
                if file_size > 0:
                    f.seek(max(0, file_size - 2000))  # 读取最后2000字符
                    lines = f.readlines()
                    if len(lines) > 1:
                        print("📋 最近的训练日志:")
                        for line in lines[-10:]:  # 显示最后10行
                            line = line.strip()
                            if line and any(keyword in line.lower() for keyword in 
                                          ['loss', 'epoch', 'step', 'lr', 'gpu']):
                                print(f"  📊 {line}")
                
                # 继续监控新内容
                while self.running:
                    line = f.readline()
                    if line:
                        line = line.strip()
                        if line:
                            # 高亮显示重要信息
                            if any(keyword in line.lower() for keyword in 
                                  ['loss', 'epoch', 'step', 'lr']):
                                print(f"📈 {line}")
                            elif any(keyword in line.lower() for keyword in 
                                    ['error', 'failed', '失败']):
                                print(f"🚨 {line}")
                            elif any(keyword in line.lower() for keyword in 
                                    ['complete', 'finished', '完成', '成功']):
                                print(f"✅ {line}")
                            else:
                                print(f"ℹ️  {line}")
                    else:
                        time.sleep(1)
                        
        except FileNotFoundError:
            logger.error(f"日志文件不存在: {latest_log}")
        except Exception as e:
            logger.error(f"监控日志时出错: {e}")
    
    def check_training_status(self):
        """检查训练状态"""
        # 检查是否有训练进程在运行
        try:
            result = subprocess.run(['pgrep', '-f', 'universal_document_trainer'], 
                                  capture_output=True, text=True)
            if result.returncode == 0:
                pids = result.stdout.strip().split('\n')
                print(f"🏃 发现训练进程: {', '.join(pids)}")
                return True
            else:
                print("💤 未发现训练进程")
                return False
        except Exception as e:
            logger.debug(f"检查进程失败: {e}")
            return False
    
    def run_interactive_monitor(self):
        """运行交互式监控"""
        print("🎯 训练监控系统启动")
        print("=" * 60)
        
        while self.running:
            try:
                # 显示菜单
                print("\n选择监控选项:")
                print("1. 监控GPU状态")
                print("2. 监控训练日志")
                print("3. 检查训练状态")
                print("4. 自动监控模式")
                print("5. 退出")
                
                choice = input("\n请选择 (1-5): ").strip()
                
                if choice == '1':
                    self.monitor_gpu_usage()
                elif choice == '2':
                    stage = input("输入阶段号 (1/2) 或回车监控所有: ").strip()
                    stage_num = int(stage) if stage and stage.isdigit() else None
                    print("🔍 开始监控训练日志 (Ctrl+C 停止)...")
                    self.monitor_training_logs(stage_num)
                elif choice == '3':
                    self.check_training_status()
                elif choice == '4':
                    self.auto_monitor()
                elif choice == '5':
                    break
                else:
                    print("❌ 无效选择")
                    
            except KeyboardInterrupt:
                print("\n⏹️  返回主菜单")
                continue
            except Exception as e:
                logger.error(f"监控出错: {e}")
    
    def auto_monitor(self):
        """自动监控模式"""
        print("🤖 启动自动监控模式...")
        print("每10秒刷新一次状态信息")
        print("按 Ctrl+C 返回主菜单")
        
        try:
            while self.running:
                os.system('clear' if os.name == 'posix' else 'cls')
                print(f"🕐 {time.strftime('%Y-%m-%d %H:%M:%S')}")
                
                # 检查训练状态
                is_training = self.check_training_status()
                
                # 显示GPU状态
                self.monitor_gpu_usage()
                
                if is_training:
                    # 显示最新日志
                    log_files = list(self.log_dir.glob("stage*/*_training.log"))
                    if log_files:
                        latest_log = max(log_files, key=lambda x: x.stat().st_mtime)
                        try:
                            with open(latest_log, 'r', encoding='utf-8') as f:
                                lines = f.readlines()
                                if lines:
                                    print("📊 最新训练信息:")
                                    for line in lines[-5:]:
                                        line = line.strip()
                                        if line and any(keyword in line.lower() for keyword in 
                                                      ['loss', 'epoch', 'step']):
                                            print(f"  {line}")
                        except:
                            pass
                
                print("\n按 Ctrl+C 返回主菜单...")
                time.sleep(10)
                
        except KeyboardInterrupt:
            print("\n⏹️  退出自动监控模式")


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='训练进度监控工具')
    parser.add_argument('--stage', type=int, choices=[1, 2], 
                       help='监控指定阶段')
    parser.add_argument('--auto', action='store_true', 
                       help='自动监控模式')
    parser.add_argument('--log-only', action='store_true', 
                       help='仅监控日志')
    
    args = parser.parse_args()
    
    monitor = TrainingMonitor()
    
    try:
        if args.log_only:
            monitor.monitor_training_logs(args.stage)
        elif args.auto:
            monitor.auto_monitor()
        else:
            monitor.run_interactive_monitor()
            
    except KeyboardInterrupt:
        print("\n👋 监控系统已停止")
    except Exception as e:
        logger.error(f"监控系统错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()